package main

import(
	"net"
	"strings"
	"fmt"
//	"crypto"
	"crypto/sha1"
	"encoding/base64"
	"io"
	"errors"
    "time"
 //   "github.com/kmalloc/gohook"
//	"encoding/hex"
)
type sockc struct{
	t int //定义客户端类型	TCP or websocket 0为TCP 1为websocket
	Conn net.Conn
	MaskingKey []byte
    init bool
    Deadtime int
    tempdata []byte
    id int64    //避免多次注册
}

func WsHand(conn *sockc,data []byte)bool{
	request:=strings.Replace(string(data), "\r\n", "\n", -1)
	k:=parseHandshake(request)
	if(strings.Contains(request, "Upgrade: websocket")){
		sd:=fmt.Sprintf("HTTP/1.1 101 Switching Protocols\r\nUpgrade: %s\r\nSec-WebSocket-Version: 13\r\nConnection: %s\r\nServer: Server: SmartH/1.0.0\r\nSec-WebSocket-Accept: %s\r\nSec-WebSocket-Origin: %s\r\nSec-WebSocket-Location: ws://%s\r\n\r\n",k["Upgrade"],k["Connection"],calc_d(k["Sec-WebSocket-Key"]),k["Origin"],k["Host"])
		conn.Write([]byte(sd))
		conn.t=1
		return true
    }
    //hook.hook()
	return false
}
func calc_d(key string)string{
	guid := "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 
        // 计算Sec-WebSocket-Accept
		h := sha1.New()
		io.WriteString(h, key + guid)
        accept := make([]byte, 28)
		base64.StdEncoding.Encode(accept, h.Sum(nil))
		return string(accept)
}
func GetBetweenStr(str, start, end string) string {
    n := strings.Index(str, start)
    if n == -1 {
        n = 0
    } else {
        n = n + len(start)  // 增加了else，不加的会把start带上
    }
    str = string([]byte(str)[n:])
    m := strings.Index(str, end)
    if m == -1 {
        m = len(str)
    }
    str = string([]byte(str)[:m])
    return str
}
func Tobase64(input []byte)string{return base64.StdEncoding.EncodeToString(input)}
func Debase64(encodeString string)[]byte{decodeBytes, _ := base64.StdEncoding.DecodeString(encodeString);return decodeBytes}
func Sha1String(data string) []byte {
	sha1 := sha1.New()
	sha1.Write([]byte(data))
	return sha1.Sum([]byte(nil))//hex.EncodeToString()
}
func parseHandshake(content string) map[string]string {
	//headers := make(map[string]string)
	var headers=make(map[string] string)
    lines := strings.Split(content, "\n")
 
    for _,line := range lines {
        if len(line) >= 0 {
			words := strings.Split(line, ": ")
            if len(words) == 2 {
				name:=words[0]
				v:=words[1]
				headers[name] = v
				
            }
        }
	}
	//print(headers["Host "]+"\n")
	//fmt.Printf("%s\n", headers["Host"])
    return headers
}
// ------------------websocket ----------------------------------
func (this *sockc)SendIframe(data []byte)(lenx int,err error)  {
    // 这里只处理data长度<125的
    x:=0
    lenth := len(data)
    //var pushHead []byte
    if lenth <126 { //>=125
        /*pushHead:=make([]byte,2);
        pushHead[0] = 0x81;
        pushHead[1]=(byte)(lenth)*/

        maskedData := make([]byte, lenth)
        for i := 0; i < lenth; i++ {
            if this.MaskingKey != nil {
               maskedData[i] = data[i] ^ this.MaskingKey[i % 4]
         }else {
             maskedData[i] = data[i]
         }
        }
        x,err=this.Conn.Write([]byte{0x81})
        if(err!=nil){return}
        var payLenByte byte
        if this.MaskingKey != nil && len(this.MaskingKey) != 4 {
            payLenByte = byte(0x80) | byte(lenth)
            x,err=this.Conn.Write([]byte{payLenByte})
            if(err!=nil){return}
            x,err=this.Conn.Write(this.MaskingKey)
            if(err!=nil){return}
        }else {
            payLenByte = byte(0x00) | byte(lenth)
            x,err=this.Conn.Write([]byte{payLenByte})
            if(err!=nil){return}
        }
        x,err=this.Conn.Write(data)
        if(err!=nil){return}
        lenx=lenx+x
        return 
    }else if(lenth<65536 && lenth >126){
        pushHead:=make([]byte,4);
        pushHead[0]=129;
        pushHead[1]=126;
        pushHead[2]=(byte)((lenth>>8) & 0xFF);
        pushHead[3]=(byte)(lenth & 0xFF);
        /*x,err=this.Conn.Write(pushHead)
        if(err!=nil){return}
        x,err=this.Conn.Write(data)
        if(err!=nil){return}*/
        //out.write(pushHead);
        x,err=this.Conn.Write(pushHead)
    lenx=x
    if(err!=nil){return}
    x,err=this.Conn.Write(data)
    lenx=lenx+x
    if(err!=nil){return}
    return
    }else{
        pushHead:=make([]byte,10);
        pushHead[0]=129//buff[0];
        pushHead[1]=127;
        pushHead[2]=(byte)((lenth>>56) & 0xFF);
        pushHead[3]=(byte)((lenth>>48) & 0xFF);
        pushHead[4]=(byte)((lenth>>40) & 0xFF);
        pushHead[5]=(byte)((lenth>>32) & 0xFF);
        pushHead[6]=(byte)((lenth>>24) & 0xFF);
        pushHead[7]=(byte)((lenth>>16) & 0xFF);
        pushHead[8]=(byte)((lenth>>8) & 0xFF);
        pushHead[9]=(byte)(lenth & 0xFF);
        x,err=this.Conn.Write(pushHead)
    lenx=x
    if(err!=nil){return}
    x,err=this.Conn.Write(data)
    lenx=lenx+x
    if(err!=nil){return}
    return
    }
    
    /*else{
        return errors.New("send iframe data error len >125")
    }*/
    //return nil
    
    
}
func (this *sockc)ReadIframe() (data []byte, err error){
    err = nil
    //第一个字节：FIN + RSV1-3 + OPCODE
    //print("enterxxxxxxxxxxx")
    opcodeByte := make([]byte, 1)
    _,err=this.Conn.Read(opcodeByte)
    if(err!=nil){return nil,err}
	FIN := opcodeByte[0] >> 7
	/*
    RSV1 := opcodeByte[0] >> 6 & 1
    RSV2 := opcodeByte[0] >> 5 & 1
    RSV3 := opcodeByte[0] >> 4 & 1
	OPCODE := opcodeByte[0] & 15
	*/
 //   log.Println(RSV1,RSV2,RSV3,OPCODE)
 
    payloadLenByte := make([]byte, 1)
    _,err=this.Conn.Read(payloadLenByte)
    if(err!=nil){return nil,err}
    payloadLen := int(payloadLenByte[0] & 0x7F)
    mask := payloadLenByte[0] >> 7
 
    if payloadLen == 127 {
        extendedByte := make([]byte, 8)
        _,err=this.Conn.Read(extendedByte)
        if(err!=nil){return nil,err}
    }
     
    maskingByte := make([]byte, 4)
    if mask == 1 {
        _,err=this.Conn.Read(maskingByte)
        if(err!=nil){return nil,err}
        this.MaskingKey = maskingByte
    }
 
    payloadDataByte := make([]byte, payloadLen)
    _,err=this.Conn.Read(payloadDataByte)
    if(err!=nil){return nil,err}
 //   log.Println("data:", payloadDataByte)
 
    dataByte := make([]byte, payloadLen)
    for i := 0; i < payloadLen; i++ {
        if mask == 1 {
            dataByte[i] = payloadDataByte[i] ^ maskingByte[i % 4]
        }else {
            dataByte[i] = payloadDataByte[i]
        }
    }
 
    if FIN == 1 {
        data = dataByte
        return
    }
    var nextData []byte
    
    //var err error 

    Try(func(){nextData, err = this.ReadIframe()},func(a interface{}){err=errors.New("close ")})
    
    
    if err != nil {
        return
    }
    data = append(data, nextData...)
    return
}
//--------------------end ----------------------------------------
func (this *sockc)Read(x []byte)(int,error){
	if(this.t==0){/*buf := make([]byte, 1024);*/return this.Conn.Read(x)}else if(this.t==1){
        g,err:=this.ReadIframe();
        k:=0
        if(len(g)>len(x)){g=g[:len(x)]}
        if(len(g)<len(x)){k=len(g)}else{k=len(x)}
        for i:=0;i<k;i++{
            x[i]=g[i]
        }
        ;return len(g),err}else{return 0,nil}
}
func (this *sockc)Write(data []byte)(i int,e error){
    try(func(){
        if(this.t==0){i,e= this.Conn.Write(data)
            }else{i,e=this.SendIframe(data)}
    })
    return
}

func makesockc(sock net.Conn,lx int)*sockc{
	//c:=make(*sockc) 
	c:=new(sockc)
	//var c sockc
	c.Conn=sock
	c.t=lx
	return c
}
func (this *sockc)SetReadDeadline(t time.Time)error{
    //if(this.t==1){gohook.Hook(fmt.Println, myPrintln, myPrintlnTramp)}
	return this.Conn.SetReadDeadline(t)
}
func (this *sockc)Close()(err error){
        for name,k:= range ul{
            if(k==this){delete(ul,name);break}
        }
    
     try(func(){ err=this.Conn.Close()})
     return
}
func (this *sockc)RemoteAddr()net.Addr{
	return this.Conn.RemoteAddr()
}
func (this *sockc)Send(lx byte,data []byte)(int,error){
    return this.Write(byte_add_to_bytes([]byte{lx},data))
}